package com.dyava.commons.token_bucket.redis;

import com.dyava.commons.token_bucket.BToken;
import com.dyava.commons.token_bucket.TokenBucket;
import org.springframework.data.redis.core.StringRedisTemplate;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.TimeUnit;

public class RedisTokenBucket extends TokenBucket {
    public static final String REDIS_PRE = "TBUCKET:";

    StringRedisTemplate stringRedisTemplate;

    public StringRedisTemplate getStringRedisTemplate() {
        return stringRedisTemplate;
    }

    public void setStringRedisTemplate(StringRedisTemplate stringRedisTemplate) {
        this.stringRedisTemplate = stringRedisTemplate;
    }

    public RedisTokenBucket(String name, int limitNum, int limitDen){
        super(name,limitNum,limitDen);
    }


    @Override
    public boolean limit(Object id) {
            String key = new StringBuilder(REDIS_PRE).append(name).append(":").append(id==null? NULL : id.toString()).toString();
            Map<Object, Object> jdata = stringRedisTemplate.opsForHash().entries(key);
            BToken token = new BToken();
            if(jdata==null || jdata.isEmpty()){
                token.setNum(limitNum);
            }else{
                token.setNum(Integer.parseInt(jdata.get("num").toString()));
                token.setLastExecuteTime(Long.parseLong(jdata.get("lastExecuteTime").toString()));
            }
            int expectNum = calculateExpectNum(token);
            if(expectNum > 0){
                expectNum--;
                token.setLastExecuteTime(System.currentTimeMillis());
                token.setNum(expectNum);
                Map<String,String> jset = new HashMap<>();
                jset.put("num",String.valueOf(token.getNum()));
                jset.put("lastExecuteTime",String.valueOf(token.getLastExecuteTime()));
                stringRedisTemplate.opsForHash().putAll(key,jset);
                stringRedisTemplate.expire(key,limitDen*2,TimeUnit.SECONDS);
                return true;
            }else{
                return false;
            }
    }

}
